package org.example.demo1.websocket;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import org.example.demo1.service.SSHTerminalService;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.*;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.UUID;

@Component
@Slf4j
public class SSHTerminalWebSocketHandler implements WebSocketHandler {

    @Autowired
    private SSHTerminalService sshTerminalService;

    private final ObjectMapper objectMapper = new ObjectMapper();
    // 存储WebSocketSession ID与SSH连接Session ID的映射
    private final Map<String, String> sessionIdMap = new ConcurrentHashMap<>();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        String sessionId = session.getId();
        log.info("WebSocket连接已建立: {}", sessionId);

        // 建议：将此消息也封装为JSON格式，与前端保持一致
        session.sendMessage(new TextMessage(objectMapper.writeValueAsString(
                Map.of("type", "websocket_established", "message", "WebSocket连接已建立"))));
    }

    @Override
    public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) throws Exception {
        try {
            String sessionId = session.getId();
            String payload = message.getPayload().toString();

            log.info("收到WebSocket消息: {} - {}", sessionId, payload);

            JsonNode jsonNode = objectMapper.readTree(payload);
            String type = jsonNode.has("type") ? jsonNode.get("type").asText() : null;

            if (type == null) {
                sendErrorMessage(session, "消息类型 (type) 缺失.");
                return;
            }

            switch (type) {
                case "connect":
                    handleConnect(session, jsonNode);
                    break;
                case "command":
                    handleCommand(session, jsonNode);
                    break;
                case "resize":
                    handleResize(session, jsonNode);
                    break;
                case "disconnect":
                    handleDisconnect(session);
                    break;
                case "ping":
                    handlePing(session); // 处理前端发送的ping消息
                    break;
                default:
                    sendErrorMessage(session, "未知消息类型: " + type);
                    break;
            }
        } catch (Exception e) {
            log.error("处理WebSocket消息失败: {}", e.getMessage(), e);
            sendErrorMessage(session, "处理消息时发生内部错误: " + e.getMessage());
            cleanupSession(session);
        }
    }

    private void handleConnect(WebSocketSession session, JsonNode jsonNode) throws Exception {
        String webSocketSessionId = session.getId();
        log.info("处理连接请求: {}", webSocketSessionId);

        JsonNode configNode = jsonNode.get("config");
        if (configNode == null || !configNode.isObject()) {
            log.error("连接请求中缺少'config'对象或'config'不是一个有效的JSON对象: {}", jsonNode.toPrettyString());
            sendErrorMessage(session, "连接请求中缺少'config'配置信息。");
            return;
        }

        String host = null;
        int port = 0;
        String username = null;
        String password = null;

        if (configNode.has("host") && configNode.get("host").isTextual()) {
            host = configNode.get("host").asText();
        } else {
            sendErrorMessage(session, "连接配置中缺少 'host' 或其格式不正确.");
            return;
        }

        if (configNode.has("port") && configNode.get("port").isIntegralNumber()) {
            port = configNode.get("port").asInt();
        } else {
            sendErrorMessage(session, "连接配置中缺少 'port' 或其格式不正确.");
            return;
        }

        if (configNode.has("username") && configNode.get("username").isTextual()) {
            username = configNode.get("username").asText();
        } else {
            sendErrorMessage(session, "连接配置中缺少 'username' 或其格式不正确.");
            return;
        }

        if (configNode.has("password")) {
            if (configNode.get("password").isTextual()) {
                password = configNode.get("password").asText();
            } else {
                sendErrorMessage(session, "连接配置中 'password' 格式不正确.");
                return;
            }
        } else {
            password = "";
        }

        String sshConnectionId = UUID.randomUUID().toString(); // 生成新的SSH会话ID

        // 调用服务层连接SSH，并获取返回的SSH连接ID
        String actualSshConnectionId = sshTerminalService.connectSSH(
                sshConnectionId, host, port, username, password, session
        );

        if (actualSshConnectionId != null) { // 如果实际连接成功并返回了ID
            sessionIdMap.put(webSocketSessionId, actualSshConnectionId); // 存储映射关系
            // ！！！重要：在"connected"消息中包含sessionId
            session.sendMessage(new TextMessage(objectMapper.writeValueAsString(
                    Map.of("type", "connected", "message", "SSH连接成功！", "sessionId", actualSshConnectionId))));
            log.info("SSH连接成功并绑定到WebSocket会话: {}", webSocketSessionId);
        } else {
            session.sendMessage(new TextMessage(objectMapper.writeValueAsString(
                    Map.of("type", "error", "message", "SSH连接失败，请检查配置或网络。"))));
            log.error("SSH连接失败 for WebSocket会话: {}", webSocketSessionId);
            sessionIdMap.remove(webSocketSessionId); // 清理map中的无效条目
        }
    }

    private void handleCommand(WebSocketSession session, JsonNode jsonNode) throws Exception {
        String webSocketSessionId = session.getId();
        String sshConnectionId = sessionIdMap.get(webSocketSessionId);

        if (sshConnectionId == null) {
            sendErrorMessage(session, "SSH连接未建立，无法发送命令。");
            return;
        }

        // --- 核心修改：将 'command' 字段改为 'data' 字段 ---
        if (!jsonNode.has("data") || !jsonNode.get("data").isTextual()) {
            sendErrorMessage(session, "命令消息中缺少 'data' 字段或其格式不正确."); // 更新错误消息
            return;
        }

        String command = jsonNode.get("data").asText(); // 从 'data' 字段获取命令
        log.info("SSH连接ID: {} - 收到命令: {}", sshConnectionId, command.replace("\n", "\\n"));
        sshTerminalService.sendCommand(sshConnectionId, command);
    }


    private void handleResize(WebSocketSession session, JsonNode jsonNode) throws Exception {
        String webSocketSessionId = session.getId();
        String sshConnectionIdForResize = sessionIdMap.get(webSocketSessionId);

        if (sshConnectionIdForResize == null) {
            sendErrorMessage(session, "SSH连接未建立，无法调整终端大小。");
            return;
        }

        int cols = 0;
        int rows = 0;
        int pixelWidth = 0;
        int pixelHeight = 0;

        if (jsonNode.has("cols") && jsonNode.get("cols").isIntegralNumber()) {
            cols = jsonNode.get("cols").asInt();
        } else {
            sendErrorMessage(session, "调整大小消息中缺少 'cols' 或其格式不正确.");
            return;
        }

        if (jsonNode.has("rows") && jsonNode.get("rows").isIntegralNumber()) {
            rows = jsonNode.get("rows").asInt();
        } else {
            sendErrorMessage(session, "调整大小消息中缺少 'rows' 或其格式不正确.");
            return;
        }

        // pixelWidth和pixelHeight是可选的，如果存在则进行解析
        if (jsonNode.has("pixelWidth") && jsonNode.get("pixelWidth").isIntegralNumber()) {
            pixelWidth = jsonNode.get("pixelWidth").asInt();
        }
        if (jsonNode.has("pixelHeight") && jsonNode.get("pixelHeight").isIntegralNumber()) {
            pixelHeight = jsonNode.get("pixelHeight").asInt();
        }

        log.info("调整终端大小SSH Connection ID: {}, cols: {}, rows: {}, pxWidth: {}, pxHeight: {}",
                sshConnectionIdForResize, cols, rows, pixelWidth, pixelHeight);
        sshTerminalService.resizeTerminal(sshConnectionIdForResize, cols, rows, pixelWidth, pixelHeight);
    }

    private void handleDisconnect(WebSocketSession session) throws Exception {
        String webSocketSessionId = session.getId();
        log.info("收到断开连接请求: {}", webSocketSessionId);

        // 修正：使用sshConnectionId来断开SSH连接
        String sshConnectionId = sessionIdMap.remove(webSocketSessionId);
        if (sshConnectionId != null) {
            sshTerminalService.disconnectSSH(sshConnectionId);
            log.info("已通过WebSocket断开请求清理SSH连接: {}", sshConnectionId);
        } else {
            log.warn("尝试断开不存在的SSH连接，WebSocket会话ID: {}", webSocketSessionId);
        }
        session.close(CloseStatus.NORMAL); // 关闭WebSocket会话
    }

    private void handlePing(WebSocketSession session) throws Exception {
        // 响应心跳
        session.sendMessage(new TextMessage(objectMapper.writeValueAsString(
                Map.of("type", "pong", "timestamp", System.currentTimeMillis())
        )));
        log.debug("向前端WebSocket会话 {} 发送心跳pong.", session.getId());
    }

    // 辅助方法：发送错误消息到WebSocket会话
    private void sendErrorMessage(WebSocketSession session, String message) throws IOException {
        session.sendMessage(new TextMessage(objectMapper.writeValueAsString(
                Map.of("type", "error", "message", message))));
        log.error("发送错误消息到WebSocket会话 {}: {}", session.getId(), message);
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        log.error("WebSocket传输错误: {} - {}\n{}", session.getId(), exception.getMessage(), exception.toString());
        sendErrorMessage(session, "WebSocket传输错误: " + exception.getMessage());
        cleanupSession(session);
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) throws Exception {
        log.info("WebSocket连接已关闭: {} - {}\n{}", session.getId(), closeStatus.getCode(), closeStatus.getReason());
        cleanupSession(session);
    }

    private void cleanupSession(WebSocketSession session) {
        String webSocketSessionId = session.getId();
        // 修正：从map中移除并获取对应的SSH连接ID
        String sshConnectionId = sessionIdMap.remove(webSocketSessionId);
        if (sshConnectionId != null) {
            sshTerminalService.disconnectSSH(sshConnectionId);
            log.info("WebSocket会话关闭时清理SSH连接: {}", sshConnectionId);
        } else {
            log.debug("没有找到对应的SSH连接来清理，WebSocket会话ID: {}", webSocketSessionId);
        }
    }

    @Override
    public boolean supportsPartialMessages() {
        return false;
    }
}